import logging
import time

import pythoncom

from lib import Config
from lib.FireWall import add_rule, RuleItem, change_rule_addr, del_rule_by_group, rename_group
from lib.SystemIP import get_connected_ips, merge_ipv4

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s - %(levelname)s - %(message)s',
                    filename='firewall.log')


class FireWallManager:
    def __init__(self):
        # { rule_name:RuleItem}
        self.rules = {}
        # 根据ip地址段划分
        # { "192.168": [1,2,3,4,5] }
        self.ipv4_map = {}
        # { "2001:0db8:85a3:0000:0000:8a2e:0370:7334": True }
        self.ipv6_map = {}
        self.group_name = Config.FireRuleName
        # 用10天的时间做索引，防止重复
        self.rule_idx = int(time.time()) % 864000

    # new rule
    def new_rule(self):
        self.rule_idx += 1

        rule_name = str(Config.FireRuleName) + "_" + str(self.rule_idx)
        rule_item = RuleItem(rule_name)
        logging.info("create new rule:%s", rule_item.name)
        return rule_item



    def init_established_ips(self):
        pythoncom.CoInitialize()
        # 清理历时规则
        self.rules = {}
        ipmap = get_connected_ips()
        for ip in ipmap['ipv4']:
            self.add_ip(ip)
        for ip in ipmap['ipv6']:
            self.ipv6_map[ip] = True

        rule_item = self.new_rule()

        for key in self.ipv4_map.keys():
            ip_arr = self.ipv4_map[key]
            rule_item.ip_str += merge_ipv4(key, ip_arr)
            rule_item.keys.append(key)
            # 一个规则最多添加10个IP
            if len(rule_item.ip_str) > 300:
                rule_item.ip_str = rule_item.ip_str[:-1]
                self.rules[rule_item.name] = rule_item
                rule_item = self.new_rule()
        if len(rule_item.keys) > 0:
            rule_item.ip_str = rule_item.ip_str[:-1]
            self.rules[rule_item.name] = rule_item

        tmp_group_name = 'tmp_group_name_xxxxx'
        # 添加规则
        for rule_name in self.rules.keys():
            rule_item = self.rules[rule_name]
            dic = {
                'Name': rule_item.name,
                'RemoteAddresses': rule_item.ip_str,
                'Grouping': tmp_group_name
            }
            add_rule(dic)
        # 删除旧的分组
        del_rule_by_group(self.group_name)
        # 重命名分组
        rename_group(tmp_group_name, self.group_name)

    # 添加IP 到分组中，如过分组包含ip，则不添加,返回False,否则返回分组key
    def add_ip(self, ip):
        # 获取IP主机位
        ip_host1, ip_host2, ip_host3, ig = ip.split(".")
        # 生成IP主机位
        key = ip_host1 + "." + ip_host2
        if key not in self.ipv4_map:
            self.ipv4_map[key] = []
        # 添加IP
        arr = self.ipv4_map[key]
        if ip_host3 in arr:
            logging.info("IP already exists %s", ip)
            return False
        arr.append(ip_host3)
        return key

    def add_ip_rule(self, ip):
        key = self.add_ip(ip)
        if key:
            # 重新生成规则
            rule_item = None
            # 是否是新IP
            is_new = True
            for rule_name in self.rules:
                rule_item = self.rules[rule_name]
                if key in rule_item.keys:
                    is_new = False
                    break
            if is_new:
                # 先看看是否有空的规则
                is_new = not rule_item or len(rule_item.keys) > 10
                logging.info("can't find matched rule_item")
                if is_new:
                    rule_item = self.new_rule()
                    self.rules[rule_item.name] = rule_item
                rule_item.keys.append(key)

            ip_str = ''
            for key in rule_item.keys:
                ip_arr = self.ipv4_map[key]
                ip_str += merge_ipv4(key, ip_arr)
            rule_item.ip_str = ip_str
            if is_new:
                dic = {
                    'Name': rule_item.name,
                    'RemoteAddresses': rule_item.ip_str,
                    'Grouping': self.group_name
                }
                add_rule(dic)
            else:
                change_rule_addr(rule_item.name, ip_str)

    # 清理规则
    def reset_rules(self):
        pass


if __name__ == '__main__':
    fw = FireWallManager()
    fw.init_established_ips()
    print(fw.rules)
    print(fw.ipv4_map)
    print(fw.ipv6_map)
    fw.add_ip_rule("192.168.33.43")
    fw.add_ip_rule("192.168.22.51")
    fw.add_ip_rule("192.168.23.51")
    fw.add_ip_rule("192.168.24.51")
    fw.add_ip_rule("192.168.25.51")

    fw.add_ip_rule("172.168.22.51")

    print(fw.rules)
    print(fw.ipv4_map)
    print(fw.ipv6_map)
